package com.nageoffer.onecoupon.framework.idempotent;

import cn.hutool.core.lang.Singleton;
import com.nageoffer.onecoupon.framework.exception.ClientException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 限流拦截器
 *
 * @author huangzhihao
 **/
@Slf4j
@Component
@Aspect
@RequiredArgsConstructor
public class RateLimitAspect {
    private final StringRedisTemplate stringRedisTemplate;
    private final static String RATE_LIMIT_LUA_PATH = "lua/rate_limit.lua";
    private final static String RATE_LIMIT_KEY_PREFIX = "rate_limit:%s_%s";

    /**
     * 获取限流策略，持久化db->redis
     *
     * @param strategyName 策略名称
     */
    private RateLimitStrategy getStrategy(String strategyName) {
        RateLimitStrategy strategy1 = new RateLimitStrategy("limitLow", 1, 10, "活动火爆稍后再试");
        RateLimitStrategy strategy2 = new RateLimitStrategy("limitMid", 3, 10, "活动火爆稍后再试");
        List<RateLimitStrategy> strategyList = new ArrayList<>();
        strategyList.add(strategy1);
        strategyList.add(strategy2);
        Map<String, RateLimitStrategy> strategyMap = strategyList.stream().collect(Collectors.toMap(RateLimitStrategy::getStrategyName, strategy -> strategy));
        return strategyMap.get(strategyName);
    }

    @Pointcut("@annotation(com.nageoffer.onecoupon.framework.idempotent.RateLimitAnnotation)")
    public void apiLogApoCut() {
    }

    @Around("apiLogApoCut()")
    public Object around(ProceedingJoinPoint joinPoint) {
        Object result = null;
        // 拿到请求方法
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method signatureMethod = signature.getMethod();
        String signatureName = signature.getName();
        // 拿到方法限流注解上的参数
        RateLimitAnnotation annotation = signatureMethod.getAnnotation(RateLimitAnnotation.class);
        if (annotation != null) {
            // 获取注解上的策略
            String strategy = annotation.strategy();
            // 获取限流策略
            RateLimitStrategy limitStrategy = getStrategy(strategy);
            if (limitStrategy != null) {
                // 获取LUA脚本并保存到Hutool单例容器中
                DefaultRedisScript<Long> luaScript = Singleton.get(RATE_LIMIT_LUA_PATH, () -> {
                    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
                    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(RATE_LIMIT_LUA_PATH)));
                    redisScript.setResultType(Long.class);
                    return redisScript;
                });
                // key为限流策略_当前请求的方法名
                List<String> keys = Collections.singletonList(String.format(RATE_LIMIT_KEY_PREFIX, strategy, signatureName));
                // 执行lua脚本
                Long count = stringRedisTemplate.execute(luaScript,
                        keys,
                        String.valueOf(limitStrategy.getLimitCount()),
                        String.valueOf(limitStrategy.getLimitTime()));
                // 判断返回结果
                if (count != null && count == 0) {
                    log.info("{}已被限流拦截", strategy);
                    throw new ClientException(limitStrategy.getFailMessage());
                }
            }
            try {
                result = joinPoint.proceed();//放行
            } catch (Throwable e) {
                throw new RuntimeException(e);
            }
        }
        return result;
    }

}